feat: add per-request context to the /v1/messages loop (#137) - #249
feat: add per-request context to the /v1/messages loop (#137)#249Zheng-Lu wants to merge 1 commit into
Conversation
Closes vllm-project#137. Follow-up to vllm-project#131, addressing the code review concern that the Messages loop took a bare `serde_json::Value` at a public API boundary while the handler separately parsed a `MessagesRequest` for typed field access. Introduces `MessagesRequestContext { typed: MessagesRequest, raw: Value }`, built once per request, and threads it through `run_messages_loop` and `run_messages_stream` in place of the untyped body. The raw body stays the thing forwarded upstream, and is deliberately not re-serialized from the typed view. `ContentBlock` catches unmodeled block types in `#[serde(other)] Unknown` and several variants model only the fields the gateway reads, so a typed round-trip would drop `cache_control` and `is_error` and collapse `image`/`redacted_thinking` into a literal `{"type":"unknown"}` block. - The typed view exposes only `tools()`, `stream()` and `model()` — the fields the loops read and never mutate. `messages` and `system` are unreachable through it, so a stale typed view cannot be read back after a round is appended. - The context owns every mutation to the upstream body (`force_stream`, `append_round`) plus the native web-search budget, replacing the two duplicated `append_round_to_history` copies with one method. - The two views diverge only where the gateway rewrites for upstream: `normalize_native_web_search` rewrites the native `web_search_20250305` declaration in the raw body, while the typed view keeps the client's original so the tool seam can still classify it. - `MessagesRequestContext::new` reuses the handler's routing parse, so neither view is built twice and a proxied request never pays for a raw parse it does not use. Native web-search validation folds into the constructor, still ahead of a streaming response committing its status, and `validate_native_web_search_request` drops off the public API. Also removes the single-field `ResolvedCall`/`ResolvedStreamCall` wrappers, which every caller immediately unwrapped, and a redundant deep clone of the assistant turn per round in the non-streaming loop. Note this keeps routing on the full `MessagesRequest` parse, so a body that declares a gateway tool but fails that parse on an unrelated field still falls through to the proxy silently. That behaviour is unchanged here and is worth a separate issue. Test Plan: - `cargo fmt -- --check`, `cargo clippy --all-targets -- -D warnings` and `cargo test --workspace` all clean: 889 passed, 0 failed, cassette replays included. - New `messages_loop_preserves_unmodeled_blocks_and_tool_result_fields_across_rounds` covers `image`, `redacted_thinking`, `citations` and a `tool_result` carrying `is_error` and `cache_control`, asserting they reach upstream unchanged across a gateway round. - Negative control: pointing `upstream_body` at the typed view fails 8 tests, including the existing Claude Code cache_control replays, which confirms the guard is not vacuous. Signed-off-by: Zheng Lu <Lz429671594@gmail.com>
Zheng-Lu
changed the title
franciscojavierarceo
left a comment
franciscojavierarceo
left a comment
There was a problem hiding this comment.
the handler builds both views from the same request, and the loops keep forwarding the raw body when appending tool rounds. i didn't find a behavior regression in the refactor. this was a source review; i haven't independently run the tests.
| gateway_map: &tool_seam::GatewayToolMap, | ||
| allowed_searches: usize, | ||
| ) -> Vec<ResolvedStreamCall> { | ||
| ) -> Vec<Value> { |
There was a problem hiding this comment.
Can we return Vec<GatewayToolResult> with typed tool_use_id, content, and is_error fields? These blocks are generated by the gateway and have a known schema, so Vec<Value> loses compile-time guarantees unnecessarily. Please carry that type through both execution helpers and append_round, converting to JSON when assembling the upstream body. The existing ContentBlock::ToolResult would need extending before reuse because it lacks is_error.
| pub struct MessagesRequestContext { | ||
| /// The client's request as received. Read-only: routing and registry | ||
| /// construction only. | ||
| typed: MessagesRequest, |
There was a problem hiding this comment.
Can we move only tools, stream, and model into the context’s typed state? Keeping the entire MessagesRequest also retains its owned message history and system prompt, although those fields are deliberately inaccessible. Previously the streaming handler dropped that parsed request after returning the response; now ctx carries this duplicate prompt data throughout the stream. For long prompts and concurrent streams, this increases retained memory without providing additional functionality.
| /// native web-search declaration. | ||
| pub fn new(typed: MessagesRequest, body: &[u8]) -> ExecutorResult<Self> { | ||
| let raw = serde_json::from_slice(body).map_err(ExecutorError::JsonError)?; | ||
| Self::from_parts(typed, raw) |
There was a problem hiding this comment.
Can construction guarantee that the typed and raw views come from the same input? new accepts them independently and only checks that body parses as JSON. A valid MessagesRequest paired with b"[]" currently returns Ok, and mismatched model/stream values are accepted too. The current HTTP handler pairs them correctly, but the public constructor does not enforce the invariant the loops rely on. Deriving both views internally, or accepting a validated paired input, would make that guarantee enforceable.
3 participants
Summary
Closes #137. Follow-up to #131, addressing the code review concern that the Messages loop took a bare
serde_json::Valueat a public API boundary while the handler separately parsed aMessagesRequestfor typed field access.Technical Details
Introduces
MessagesRequestContext { typed: MessagesRequest, raw: Value }, built once per request, and threads it throughrun_messages_loopandrun_messages_streamin place of the untyped body.The raw body stays the thing forwarded upstream, and is deliberately not re-serialized from the typed view.
ContentBlockcatches unmodeled block types in#[serde(other)] Unknownand several variants model only the fields the gateway reads, so a typed round-trip would dropcache_controlandis_errorand collapseimage/redacted_thinkinginto a literal{"type":"unknown"}block.tools(),stream()andmodel()— the fields the loops read and never mutate.messagesandsystemare unreachable through it, so a stale typed view cannot be read back after a round is appended.force_stream,append_round) plus the native web-search budget, replacing the two duplicatedappend_round_to_historycopies with one method.normalize_native_web_searchrewrites the nativeweb_search_20250305declaration in the raw body, while the typed view keeps the client's original so the tool seam can still classify it.MessagesRequestContext::newreuses the handler's routing parse, so neither view is built twice and a proxied request never pays for a raw parse it does not use. Native web-search validation folds into the constructor, still ahead of a streaming response committing its status, andvalidate_native_web_search_requestdrops off the public API.Also removes the single-field
ResolvedCall/ResolvedStreamCallwrappers, which every caller immediately unwrapped, and a redundant deep clone of the assistant turn per round in the non-streaming loop.Note this keeps routing on the full
MessagesRequestparse, so a body that declares a gateway tool but fails that parse on an unrelated field still falls through to the proxy silently. That behaviour is unchanged here and is worth a separate issue.Test Plan
cargo fmt -- --check,cargo clippy --all-targets -- -D warningsandcargo test --workspaceall clean: 889 passed, 0 failed, cassette replays included.messages_loop_preserves_unmodeled_blocks_and_tool_result_fields_across_roundscoversimage,redacted_thinking,citationsand atool_resultcarryingis_errorandcache_control, asserting they reach upstream unchanged across a gateway round.upstream_bodyat the typed view fails 8 tests, including the existing Claude Code cache_control replays, which confirms the guard is not vacuous.